import os
import argparse
from tkinter import Tk
from tkinter.filedialog import askopenfilename

os.environ['CUDA_VISIBLE_DEVICES']='1'
import torch
import torchvision.transforms as transforms
from torchvision import models
from PIL import Image
import numpy as np
from scipy.spatial import distance
import torch.nn as nn
import pdb


class SiameseNetwork(nn.Module):
    def __init__(self):
        super(SiameseNetwork, self).__init__()
        self.vgg = models.vgg16(pretrained=False)
        # 删除最后的全连接层
        self.vgg.classifier = self.vgg.classifier[:-1]
        
    def forward(self, img1):
        out1 = self.vgg(img1)
        return out1
    
# 解析命令行参数
parser = argparse.ArgumentParser(description="Image retrieval system.")
parser.add_argument("--use_pretrained", default='True', type=str, help="Use pretrained model or load checkpoint.")
parser.add_argument("--checkpoint_path", default='', type=str, help="Path to checkpoint.")
parser.add_argument("--target_image", required=False, type=str, help="Path to the target image.")
args = parser.parse_args()

# 加载模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# pdb.set_trace()
if args.use_pretrained.lower() == 'true':
    model = models.vgg16(pretrained=True)
else:
    model = SiameseNetwork()
    weights = torch.load(args.checkpoint_path)
    model.load_state_dict(weights)
    # weights = torch.load(args.checkpoint_path)
    # model.load_state_dict(weights)


model.eval()

# 定义图像转换器
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 提取特征并保存
def extract_features(parent_folder):
    # 图像检索系统的实现函数一，使用神经网络模型将图像数据库的图像提取出来，存储为npy文件
    # pdb.set_trace()
    features_dict = {}

    for child_folder in os.listdir(parent_folder):
        child_folder_path = os.path.join(parent_folder, child_folder)
        for image_name in os.listdir(child_folder_path):
            image_path = os.path.join(child_folder_path, image_name)
            image = Image.open(image_path)
            image = transform(image).unsqueeze(0)
            if args.use_pretrained.lower() == 'true':
                features = model.features(image).detach().numpy().flatten()
            else:
                features = model.vgg.features(image).detach().numpy().flatten()

            features_dict[image_name] = features

    # 保存为二进制文件
    if args.use_pretrained.lower() == 'true':
        np.save('特征向量库/features_pretrain.npy', features_dict)
    else:
        np.save('特征向量库/features_30.npy', features_dict)
    


def search_similar(target_image):
    # 图像检索系统的实现函数二，用于将用户输入照片，与特征向量库的图像做相似度计算，并返回结果
    target_image = Image.open(target_image)
    target_image = transform(target_image).unsqueeze(0)
    if args.use_pretrained.lower() == 'true':
        target_features = model.features(target_image).detach().numpy().flatten()
        features_dict = np.load('特征向量库/features_pretrain.npy', allow_pickle=True).item()
    else:
        target_features = model.vgg.features(target_image).detach().numpy().flatten()
        features_dict = np.load('特征向量库/features_30.npy', allow_pickle=True).item()

    result = {}
    for image_name, features in features_dict.items():
        sim = 1 - distance.cosine(target_features, features)
        result[image_name] = sim

    # 按相似度排序
    sorted_result = {k: v for k, v in sorted(result.items(), key=lambda item: item[1], reverse=True)}
    # for i,j in sorted_result.items():
        # print(i,j)
    for i, j in sorted_result.items():
        print(f'{i:<30} {j}')
    return sorted_result

def compare_features(path1,path2):
    """
    函数功能:用于比较图数据库中特征信息的npy中不同图像的权重是否相同,目的是检测经过微调的模型的权重是否更新
    path1,path2分别是npy文件的位置
    """
    ft_features_dict = np.load(path1, allow_pickle=True).item()
    features_dict = np.load(path2, allow_pickle=True).item()

    # 判断两个字典的键是否完全相同，如果不同，输出不同的键
    if ft_features_dict.keys() != features_dict.keys():
        print("The two dictionaries have different keys.")
        diff_keys = ft_features_dict.keys() ^ features_dict.keys() # 使用异或运算找出不同的键
        print("The different keys are: ", diff_keys)

    # 遍历字典键，对比features
    for key in ft_features_dict:
        if np.array_equal(ft_features_dict[key], features_dict[key]):
            print(f"The features of {key} are identical in both dictionaries.")
        else:
            print(f"The features of {key} are different in the dictionaries.")


if __name__=='__main__':
    # 弹出一个窗口供用户选择图片
    # Tk().withdraw() # we don't want a full GUI, so keep the root window from appearing
    # target_image = askopenfilename() # this will be the name of the image file chosen by the user
    
    # step1,取消注释，提取所有图像的特征信息
    # extract_features('newdata')

    # step2,接受用户输入，返回相似图像
    search_similar(args.target_image)

    # 比较两个模型生成的特征向量是否相同
    # compare_features('特征向量库/features_30.npy','特征向量库/features_pretrain.npy')

    


